#!/usr/bin/python3
# Prof. Shi Lei's Iterative Stepwise Optimization Algorithm
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

import numpy
import matplotlib
matplotlib.use('Agg'); # plot and save, never showing

import matplotlib.pyplot as plt

def calcXvector(clussprs):
  '''
# for a paper p:
#  x(p)=1/sqrt(|#c(p)|)
# input:
#  clussprs: cluster seperators: [0,p_1,p_2,...,p_n,|#P|]
  '''
  xvec=numpy.zeros(clussprs[-1])+1.;
  for c in range(1,len(clussprs)):
   xvec[range(clussprs[c-1],clussprs[c])]=1./numpy.sqrt(clussprs[c]-clussprs[c-1]);
  return xvec;

def monitorXandErr(eigvec,clussprs):
  '''
# monitoring x vector and error function
  '''
  xvec=calcXvector(clussprs);
  reerr=1-xvec.dot(eigvec)**2/xvec.dot(xvec)/eigvec.dot(eigvec);
  return xvec,reerr;

def visualizeDivide(eigvec,clussprs):
  '''
# visualize dividing progress
  '''
  mappedxvec=calcXvector(clussprs);
  mappedeigv=eigvec;
  maponexvec=mappedxvec/numpy.sqrt(len(clussprs)-1);
  plt.clf();
  plt.plot(range(0,len(eigvec)),mappedeigv,color="green",linewidth=2,linestyle='--',label='eigenvector');
  plt.plot(range(0,len(eigvec)),maponexvec,color="blue", linewidth=2,linestyle='-',label='x vector');
  plt.grid();
  plt.xlabel('paper index (sorted by eigvec)');
  plt.ylabel('vector component');
  plt.legend(loc='upper left');
  plt.savefig('../out/xveig'+str(len(clussprs)-1)+'.png');
  # plt.show(); # don't forget to commentize "matplotlib.use('Agg');" in the beginning

def clussprs2clusters(clussprs,sortmap):
  '''
# generate clusters set from cluster seperator data
  '''
  def _idxfunc(x):
    return x[1];
  clusters=[list(set([_idxfunc(sortmap[p]) for p in range(clussprs[c],clussprs[c+1])])) for c in range(0,len(clussprs)-1)];
  return clusters;

def calcClusters(eigvec, clusnum=20, monitor=False):
  '''
# set paper clusters to make b=sum(a|c)/sqrt(#c) approching sum(a)/sum(sqrt(#c))
# Algorithm by Shi Lei <shijim@gmail.com>
# Program by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  eigenvector
#  clusnum: max number of clusters
#  monitor: whether visualize the progress
# output:
#  snapshot (last history state detail): clusters set [[paper indices],...]
  '''
  dimension=len(eigvec);

  def _keyfunc(x):
   return x[0];
  sortmap=sorted([(eigvec[p],p) for p in range(0,len(eigvec))],key=_keyfunc); # (eigen_value,origin_order)
  sortbck=sorted([(sortmap[p][1],p) for p in range(0,len(eigvec))],key=_keyfunc); # (origin_order,eigen_order)
  mappedeigv=numpy.float64([e[0] for e in sortmap]);
  numpyeigvec=numpy.array(mappedeigv); # different from list when slicing

  if (monitor):
   print('Iterative Stepwise Optimization:'+'\n  paper number: '+str(dimension)+'\n  set cluster number: '+str(clusnum));
  if (dimension<clusnum):
   print('Setting Cluster Number is smaller than Current Cluster Number, reset to equal.');
   clusnum=dimension;
  divnum=int(numpy.floor(dimension/clusnum));
  resnum=int(dimension%clusnum);
  clussprs=numpy.array(range(0,resnum)); # seperators
  clussprs*=(divnum+1);
  clussprs=list(clussprs);
  clussprs.extend(list(numpy.array(range(resnum,clusnum))*divnum+resnum));
  clussprs.append(len(numpyeigvec));
  clussprs=numpy.array(clussprs);
  eigsum=sum(numpyeigvec);
  changenum=1; # do while
  numvec=numpy.array([(clussprs[it+1]-clussprs[it]) for it in range(0,clusnum)]);
  sqrtnumvec=numpy.sqrt(numvec);
  sqrtnumsum=sum(sqrtnumvec);
  eigclussum=numpy.array([sum(numpyeigvec[clussprs[it]:clussprs[it+1]]) for it in range(0,clusnum)]); # sum eigvec value
  nextsum=0 # initialize temp var
  paralleldegree=sum(eigclussum/sqrtnumvec);
  bestrcd=paralleldegree; # best record
  noimprovecount=0; # prevent loop
  while (changenum>0 and noimprovecount<=7): # no improvement for 7 times => treat as loop
   changenum=0;
   newclussprs=[0];
   neweigclussum=[];
   for cit in range(1,clusnum):
    newclussprs.append(newclussprs[cit-1]+1);
    neweigclussum.append(sum(numpyeigvec[newclussprs[cit-1]:newclussprs[cit]]));
    if (newclussprs[cit]>dimension-clusnum+cit-1): # insufficient room for the rest clusters
     nextsum=numpy.Inf;
    else:
     nextsum=neweigclussum[cit-1]+numpyeigvec[newclussprs[cit]];
    while (nextsum/numpy.sqrt(newclussprs[cit]-newclussprs[cit-1]+1)<=eigsum/sqrtnumsum):
     neweigclussum[cit-1]+=numpyeigvec[newclussprs[cit]];
     newclussprs[cit]+=1; # move the seperator by 1
     changenum+=1;
     if (newclussprs[cit]>dimension-clusnum+cit-1): # insufficient room for the rest clusters
      nextsum=numpy.Inf;
     else:
      nextsum=neweigclussum[cit-1]+numpyeigvec[newclussprs[cit]];
   newclussprs.append(len(numpyeigvec));
   neweigclussum.append(sum(numpyeigvec[clussprs[clusnum-1]:clussprs[clusnum]]));
   clussprs=newclussprs;
   eigclussum=numpy.array(neweigclussum);
   numvec=numpy.array([(clussprs[it+1]-clussprs[it]) for it in range(0,clusnum)]);
   sqrtnumvec=numpy.sqrt(numvec);
   sqrtnumsum=sum(sqrtnumvec); # update sqrtnumsum
   paralleldegree=sum(eigclussum/sqrtnumvec);
   if (paralleldegree>bestrcd):
    bestrcd=paralleldegree;
    noimprovecount=0;
   else:
    noimprovecount+=1;
   if monitor:
     xvec,reerr=monitorXandErr(mappedeigv,list(clussprs));
     visualizeDivide(mappedeigv,list(clussprs));
     print('unparallel ratio = '+str(reerr));
     print('change num: '+str(changenum));
     print('average b objective: '+str(eigsum/sqrtnumsum));
  clusters=clussprs2clusters(list(clussprs),sortmap);
  return clusters;

